{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DINOv3 Similarity Visualization\n",
    "\n",
    "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/geoai/blob/main/docs/examples/DINOv3_visualization.ipynb)\n",
    "\n",
    "This notebook demonstrates the new visualization functionality for DINOv3 similarity analysis in the geoai package.\n",
    "\n",
    "## Install packages\n",
    "\n",
    "To use the new functionality, ensure the required packages are installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install geoai-py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import geoai\n",
    "from geoai.dinov3 import DINOv3GeoProcessor, visualize_similarity_results\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"https://huggingface.co/spaces/pszemraj/dinov3-viz-sat493m/resolve/main/aralsea_tmo_2000238_lrg.jpg\"\n",
    "test_image_path = geoai.download_file(url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create and display test image\n",
    "test_image = Image.open(test_image_path)\n",
    "\n",
    "plt.figure(figsize=(8, 8))\n",
    "plt.imshow(test_image)\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method 1: Quick Visualization with visualize_similarity_results()\n",
    "\n",
    "The easiest way to create and visualize similarity results in one step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query coordinates\n",
    "query_coords = (1250, 1250)\n",
    "\n",
    "# Side-by-side visualization\n",
    "results = visualize_similarity_results(\n",
    "    input_image=test_image_path,\n",
    "    query_coords=query_coords,\n",
    "    overlay=False,  # Side-by-side view\n",
    "    figsize=(15, 6),\n",
    "    colormap=\"turbo\",\n",
    "    show_query_point=True,\n",
    ")\n",
    "\n",
    "plt.show()\n",
    "print(f\"Query patch coordinates: {results['patch_coords']}\")\n",
    "print(f\"Patch grid size: {results['patch_grid_size']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Overlay visualization\n",
    "results_overlay = visualize_similarity_results(\n",
    "    input_image=test_image_path,\n",
    "    query_coords=query_coords,\n",
    "    overlay=True,  # Overlay view\n",
    "    alpha=0.6,  # Transparency\n",
    "    colormap=\"jet\",\n",
    "    show_query_point=True,\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method 2: Step-by-step with DINOv3GeoProcessor\n",
    "\n",
    "For more control over the process, you can use the processor methods directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize processor\n",
    "processor = DINOv3GeoProcessor()\n",
    "\n",
    "# Load and preprocess image\n",
    "data, metadata = processor.load_image(test_image_path)\n",
    "image = processor.preprocess_image_for_dinov3(data)\n",
    "features, h_patches, w_patches = processor.extract_features(image)\n",
    "\n",
    "print(f\"Image size: {image.size}\")\n",
    "print(f\"Features shape: {features.shape}\")\n",
    "print(f\"Patch grid: {h_patches} x {w_patches}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert pixel coordinates to patch coordinates\n",
    "img_w, img_h = data.shape[2], data.shape[1]\n",
    "patch_x = int((query_coords[0] / img_w) * w_patches)\n",
    "patch_y = int((query_coords[1] / img_h) * h_patches)\n",
    "\n",
    "print(f\"Query pixel coordinates: {query_coords}\")\n",
    "print(f\"Query patch coordinates: ({patch_x}, {patch_y})\")\n",
    "\n",
    "# Compute similarity\n",
    "similarities = processor.compute_patch_similarity(features, patch_x, patch_y)\n",
    "similarity_array = similarities.cpu().numpy()\n",
    "\n",
    "print(f\"Similarity range: {similarity_array.min():.3f} - {similarity_array.max():.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize patch grid\n",
    "fig = processor.visualize_patches(\n",
    "    image=image, features=features, patch_coords=(patch_x, patch_y), figsize=(12, 8)\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create side-by-side visualization\n",
    "fig = processor.visualize_similarity(\n",
    "    source=test_image_path,\n",
    "    similarity_data=similarity_array,\n",
    "    query_coords=query_coords,\n",
    "    patch_coords=(patch_x, patch_y),\n",
    "    overlay=False,\n",
    "    colormap=\"jet\",\n",
    "    figsize=(15, 6),\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create overlay visualization\n",
    "fig = processor.visualize_similarity(\n",
    "    source=test_image_path,\n",
    "    similarity_data=similarity_array,\n",
    "    query_coords=query_coords,\n",
    "    overlay=True,\n",
    "    alpha=0.5,\n",
    "    colormap=\"jet\",\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method 3: Create Overlay Image Array\n",
    "\n",
    "You can also create overlay images programmatically for further processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create overlay as numpy array\n",
    "overlay_img = processor.create_similarity_overlay(\n",
    "    source=test_image_path, similarity_data=similarity_array, colormap=\"jet\", alpha=0.6\n",
    ")\n",
    "\n",
    "# Display the overlay\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(overlay_img)\n",
    "plt.title(\"Programmatically Created Overlay\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()\n",
    "\n",
    "# Save overlay as image file\n",
    "overlay_pil = Image.fromarray((overlay_img * 255).astype(np.uint8))\n",
    "overlay_pil.save(\"demo_overlay.png\")\n",
    "print(\"Overlay saved to demo_overlay.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Different Query Points\n",
    "\n",
    "Let's test similarity analysis with different query points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test multiple query points\n",
    "query_points = [\n",
    "    (128, 128),  # Red circle\n",
    "    (384, 128),  # Green circle\n",
    "    (128, 384),  # Blue circle\n",
    "    (384, 384),  # Yellow circle\n",
    "]\n",
    "\n",
    "colors = [\"Red\", \"Green\", \"Blue\", \"Yellow\"]\n",
    "\n",
    "fig, axes = plt.subplots(2, 2, figsize=(20, 16))\n",
    "axes = axes.flatten()\n",
    "\n",
    "for i, (coords, color) in enumerate(zip(query_points, colors)):\n",
    "    # Quick visualization for each point\n",
    "    results = visualize_similarity_results(\n",
    "        input_image=test_image_path,\n",
    "        query_coords=coords,\n",
    "        overlay=True,\n",
    "        alpha=0.6,\n",
    "        show_query_point=True,\n",
    "    )\n",
    "\n",
    "    # Close the generated figure and create our own subplot\n",
    "    plt.close(results[\"visualization\"])\n",
    "\n",
    "    # Get the similarity data and create our own visualization\n",
    "    similarity_data = results[\"image_dict\"][\"image\"][0]\n",
    "    data, _ = processor.load_image(test_image_path)\n",
    "    display_img = np.transpose(data, (1, 2, 0))\n",
    "\n",
    "    # Show on subplot\n",
    "    axes[i].imshow(display_img / 255.0)\n",
    "    im = axes[i].imshow(similarity_data, cmap=\"turbo\", alpha=0.6, vmin=0, vmax=1)\n",
    "    axes[i].plot(\n",
    "        coords[0],\n",
    "        coords[1],\n",
    "        \"r*\",\n",
    "        markersize=15,\n",
    "        markeredgecolor=\"white\",\n",
    "        markeredgewidth=2,\n",
    "    )\n",
    "    axes[i].set_title(f\"Query: {color} Circle at {coords}\")\n",
    "    axes[i].axis(\"off\")\n",
    "\n",
    "    # Add colorbar\n",
    "    plt.colorbar(im, ax=axes[i], fraction=0.046, pad=0.04)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
